from functools import wraps

from webutils.ratelimit.exceptions import RateLimited
from webutils.layers.connection import ConnectionProxy
from webutils.ratelimit import is_limited
from webutils.ratelimit.services import Ratelimit


key_handlers = {
    'user': lambda ws: ws.session.user.id,
}


def ratelimit(rate, key='user', group=None, ignore=None):
    if ignore:
        assert callable(ignore)
    if key not in key_handlers:
        assert callable(key)

    def decorator(fn):
        @wraps(fn)
        async def wrapper(request_handler, *args, **kwargs):
            if ignore and ignore(request_handler, *args, **kwargs):
                return await fn(request_handler, *args, **kwargs)
            if callable(key):
                visitor_tag = key(request_handler)

            else:
                if request_handler.session.get('user'):
                    visitor_tag = key_handlers[key](request_handler)
                elif args[0].get('from'):
                    visitor_tag = args[0]['from']['id']
                else:
                    # 意外的情况
                    request_handler.logger.error(f'error:意外的情况,'
                                                 f'属于本地消息却没有from,-[{args=}]')
                    raise
            groups = [wrapper.__module__, wrapper.__name__]
            limited_resource = '.'.join(groups)
            async with ConnectionProxy(request_handler.application,
                                       request_handler.log_context) as conn:
                ratelimit_service = Ratelimit(conn)
                ratelimited = await is_limited(
                    visitor_tag,
                    rate,
                    limited_resource=limited_resource,
                    limited_group=group,
                    cache=ratelimit_service,
                )
            if ratelimited:
                raise RateLimited()
            return await fn(request_handler, *args, **kwargs)

        return wrapper

    return decorator
